
###############################################################################
### Did 3G Make Afghan Insurgents Fight More Effectively? 
### Replication files for analyses
### Mehmet Erdem Arslan
###
### Custom functions for Matched Wake Analyses 
###############################################################################


require(xtable)

# Large table for detailed results 
large_table <- function(results, caption, label) { 
  matching <- data.frame(results$matching)
  sutva <- data.frame(results$SUTVA)
  table <- cbind(results$estimates, matching, sutva)
  if (!is.null(results$parameters$estimationControls)) {
    table <- table[,-c(6:(5 + 2*length(results$parameters$estimationControls)))]
  }
  table <- table[,c(1,2,3,4,5, 10,11,12,13,20,23, 14,15,16,17,21,24)]
  # table <- table[table$pvalue < 0.1, ]
  rownames(table) <- 1:nrow(table)
  xtable(table, caption = caption, label = label, digits = 2)
}


# Custom plotting function for presenting MWA results
# The code below is taken from mwa package in R (Schutte and Donnay, 2014) and 
# modified to present results without the contour and with numbers instead 
my.filled.contour <- function (x = seq(0, 1, length.out = nrow(z)), y = seq(0, 1, 
                                                       length.out = ncol(z)), z, xlim = range(x, finite = TRUE), 
          ylim = range(y, finite = TRUE), zlim = range(z, finite = TRUE), 
          levels = pretty(zlim, nlevels), nlevels = 20, color.palette = function(n) hcl.colors(n, 
                                                                                               "YlOrRd", rev = TRUE), col = color.palette(length(levels) - 
                                                                                                                                            1), plot.title, plot.axes, key.title, key.axes, asp = NA, 
          xaxs = "i", yaxs = "i", las = 1, axes = TRUE, frame.plot = axes, 
          ...) 
{
  if (missing(z)) {
    if (!missing(x)) {
      if (is.list(x)) {
        z <- x$z
        y <- x$y
        x <- x$x
      }
      else {
        z <- x
        x <- seq.int(0, 1, length.out = nrow(z))
      }
    }
    else stop("no 'z' matrix specified")
  }
  else if (is.list(x)) {
    y <- x$y
    x <- x$x
  }
  if (any(diff(x) <= 0) || any(diff(y) <= 0)) 
    stop("increasing 'x' and 'y' values expected")
  mar.orig <- (par.orig <- par(c("mar", "las", "mfrow")))$mar
  on.exit(par(par.orig))
  w <- (3 + mar.orig[2L]) * par("csi") * 2.54
  # layout(matrix(c(2, 1), ncol = 2L), widths = c(1, lcm(w)))
  par(las = las)
  mar <- mar.orig
  mar[4L] <- mar[2L]
  mar[2L] <- 1
  par(mar = mar)
  # plot.new()
  # plot.window(xlim = c(0, 1), ylim = range(levels), xaxs = "i", 
  #             yaxs = "i")
  # rect(0, levels[-length(levels)], 1, levels[-1L], col = col)
  # if (missing(key.axes)) {
  #   if (axes) 
  #     axis(4)
  # }
  # else key.axes
  # box()
  if (!missing(key.title)) 
    key.title
  mar <- mar.orig
  mar[4L] <- 1
  par(mar = mar)
  plot.new()
  plot.window(xlim, ylim, "", xaxs = xaxs, yaxs = yaxs, asp = asp)
  .filled.contour(x, y, z, levels, col)
  if (missing(plot.axes)) {
    if (axes) {
      title(main = "", xlab = "", ylab = "")
      Axis(x, side = 1)
      Axis(y, side = 2)
    }
  }
  else plot.axes
  if (frame.plot) 
    box()
  if (missing(plot.title)) 
    title(...)
  else plot.title
  invisible()
}



my.mwa.plot <- function(x, zlim = NA, plotNAs = TRUE, ...){
    density <- 3
    lty <- 2
    lwd <- 2
    pdata <- x$estimates
    t_unit <- x$parameters$t_unit
    alpha1 <- x$parameters$alpha1
    alpha2 <- x$parameters$alpha2
    yAxis <- sort(unique(pdata$t_window))
    xAxis <- sort(unique(pdata$spat_window))
    pswcplot <- matrix(nrow=length(xAxis)+2,ncol=length(yAxis)+2)
    eswcplot <- matrix(nrow=length(xAxis)+2,ncol=length(yAxis)+2)
    for (y in 1:length(yAxis)){
      for (x in 1:length(xAxis)){
        eswcplot[x+1,y+1] <- as.numeric(pdata$estimate[pdata$t_window == yAxis[y] & pdata$spat_window == xAxis[x]])
        pswcplot[x+1,y+1] <- as.numeric(pdata$pvalue[pdata$t_window == yAxis[y] & pdata$spat_window == xAxis[x]])
      }
    }
    if (plotNAs){
      eswcplot[is.na(eswcplot)] <- 0
    }
    for (y in 1:(length(yAxis)+2)){
      eswcplot[1,y] <- eswcplot[2,y]
      eswcplot[length(xAxis)+2,y] <- eswcplot[length(xAxis)+1,y]
    }
    for (x in 1:(length(xAxis)+2)){
      eswcplot[x,1] <- eswcplot[x,2]
      eswcplot[x,length(yAxis)+2] <- eswcplot[x,length(yAxis)+1]
    }
    if (length(xAxis)==1){
      xvals <- c(0,0.5,1)
      xlims <- c(0,1)
      xat <- 0.5
      xhalfRectSize <- 1
      xstepsize <- 1 
    }else{
      xvals <- seq(-1/(length(xAxis)-1),1+1/(length(xAxis)-1),length.out=length(xAxis)+2)
      xlims <- c(-1/(length(xAxis)-1)/2,1+1/(length(xAxis)-1)/2)
      xat <- seq(0,1,by=1/(length(xAxis)-1))
      xhalfRectSize <- (1 / (length(xAxis) -1)) / 2
      xstepsize <- 1 / (length(xAxis) - 1)
    }
    if (length(yAxis)==1){
      yvals <- c(0,0.5,1)
      ylims <- c(0,1)
      yat <- 0.5
      yhalfRectSize <- 1
      ystepsize <- 1 
    }else{
      yvals <- seq(-1/(length(yAxis)-1),1+1/(length(yAxis)-1),length.out=length(yAxis)+2)
      ylims <- c(-1/(length(yAxis)-1)/2,1+1/(length(yAxis)-1)/2)
      yat <- seq(0,1,by=1/(length(yAxis)-1))
      yhalfRectSize <- (1 / (length(yAxis) -1)) / 2
      ystepsize <- 1 / (length(yAxis) - 1)
    }
    if (is.na(zlim[1])){
      zlims <- range(eswcplot[!is.na(eswcplot)], finite = TRUE)
    }else{
      zlims <- zlim
    }
    
    my.filled.contour(x = xvals, y = yvals,eswcplot,xlab = "Spatial window in kilometers",
                   ylab=paste("Temporal window in ",t_unit,sep=""),
                   # nlevels = 20,
                   nlevels = 1,
                   # color.palette = gray.colors,
                   col = "white",
                   xlim = xlims,
                   ylim = ylims,
                   zlim = zlims,
                   plot.axes = {
                     axis(1,labels=xAxis,at=xat)
                     axis(2,at=yat,labels=yAxis)
                     for (y in 1:length(yAxis)){
                       for(x in 1:length(xAxis)){
                         if (is.na(pswcplot[x+1,y+1]) | is.nan(pswcplot[x+1,y+1]) | pswcplot[x+1,y+1] <= alpha2){
                           #  rect(c(((x - 1)*xstepsize) - xhalfRectSize,((x - 1)*xstepsize) - xhalfRectSize),c(((y - 1)*ystepsize) - yhalfRectSize,((y - 1)*ystepsize) - yhalfRectSize),c(((x - 1)*xstepsize) + xhalfRectSize,((x - 1)*xstepsize) + xhalfRectSize),c(((y - 1)*ystepsize) + yhalfRectSize,((y - 1)*ystepsize)+yhalfRectSize),density=density,lwd=lwd)
                           rect(c(((x - 1)*xstepsize) - xhalfRectSize,((x - 1)*xstepsize) - xhalfRectSize),c(((y - 1)*ystepsize) - yhalfRectSize,((y - 1)*ystepsize) - yhalfRectSize),c(((x - 1)*xstepsize) + xhalfRectSize,((x - 1)*xstepsize) + xhalfRectSize),c(((y - 1)*ystepsize) + yhalfRectSize,((y - 1)*ystepsize)+yhalfRectSize), lwd=lwd, col = "lightgrey")
                           text(c((x - 1)*xstepsize,(x - 1)*xstepsize),c((y - 1)*ystepsize,(y - 1)*ystepsize), labels = paste0(as.character(round(eswcplot[x+1,y+1], digits = 2)), as.character(paste0("\n(", round(pswcplot[x+1,y+1], digits = 2), ")")) ))
                         } else {
                           rect(c(((x - 1)*xstepsize) - xhalfRectSize,((x - 1)*xstepsize) - xhalfRectSize),c(((y - 1)*ystepsize) - yhalfRectSize,((y - 1)*ystepsize) - yhalfRectSize),c(((x - 1)*xstepsize) + xhalfRectSize,((x - 1)*xstepsize) + xhalfRectSize),c(((y - 1)*ystepsize) + yhalfRectSize,((y - 1)*ystepsize)+yhalfRectSize), lwd=0.5)
                           text(c((x - 1)*xstepsize,(x - 1)*xstepsize),c((y - 1)*ystepsize,(y - 1)*ystepsize), labels = paste0(as.character(round(eswcplot[x+1,y+1], digits = 2)), as.character(paste0("\n(", round(pswcplot[x+1,y+1], digits = 2), ")")) ))
                         }
                         #  } else if (is.na(pswcplot[x+1,y+1]) | is.nan(pswcplot[x+1,y+1]) | pswcplot[x+1,y+1] > alpha1 & pswcplot[x+1,y+1] <= alpha2){
                         #  rect(c(((x - 1)*xstepsize) - xhalfRectSize,((x - 1)*xstepsize) - xhalfRectSize),c(((y - 1)*ystepsize) - yhalfRectSize,((y - 1)*ystepsize) - yhalfRectSize),c(((x - 1)*xstepsize) + xhalfRectSize,((x - 1)*xstepsize) + xhalfRectSize),c(((y - 1)*ystepsize) + yhalfRectSize,((y - 1)*ystepsize)+yhalfRectSize),density=density,lty=lty, lwd=lwd)
                         # }
                       }
                     }
                   }
    )
    
    
}
